
package com.rymcu.forest.openai;

import com.rymcu.forest.core.result.GlobalResult;
import com.rymcu.forest.core.result.GlobalResultGenerator;
import com.rymcu.forest.entity.User;
import com.rymcu.forest.openai.entity.ChatMessageModel;
import com.rymcu.forest.openai.entity.ChatModel;
import com.rymcu.forest.openai.service.OpenAiService;
import com.rymcu.forest.openai.service.SseService;
import com.rymcu.forest.rag.RAGQueryResult;
import com.rymcu.forest.rag.SuperSQLService;
import com.rymcu.forest.util.Html2TextUtil;
import com.rymcu.forest.util.UserUtils;
import com.theokanning.openai.completion.chat.ChatCompletionRequest;
import com.theokanning.openai.completion.chat.ChatMessage;
import org.apache.commons.lang.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.*;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;

/**
 * RAG增强的AI聊天控制器
 * 集成SuperSQL实现基于博客内容的智能问答
 */
@RestController
@RequestMapping("/api/v1/rag")
public class RAGChatController {

    @Autowired
    private SuperSQLService superSQLService;

    @Autowired
    private SseService sseService;

    @Value("${openai.token}")
    private String token;

    /**
     * RAG增强的聊天接口
     */
    @PostMapping("/chat")
    public GlobalResult ragChat(@RequestBody ChatModel chatModel) {
        try {
            // 获取用户信息
            User user = UserUtils.getCurrentUserByToken();

            // 获取最新的用户消息作为问题
            List<ChatMessageModel> messages = chatModel.getMessages();
            if (messages.isEmpty()) {
                return GlobalResultGenerator.genErrorResult("消息不能为空");
            }

            ChatMessageModel lastMessage = messages.get(messages.size() - 1);
            String userQuestion = lastMessage.getContent();
            String model = chatModel.getModel();

            // 1. 使用SuperSQL查询相关博客内容
            RAGQueryResult queryResult = superSQLService.queryByNaturalLanguage(userQuestion);

            // 2. 构建增强的上下文
            List<ChatMessage> enhancedMessages = buildEnhancedContext(messages, queryResult);

            // 3. 调用AI进行回答
            return sendMessageWithRAG(user, enhancedMessages, model, queryResult);

        } catch (Exception e) {
            return GlobalResultGenerator.genErrorResult("RAG聊天失败: " + e.getMessage());
        }
    }

    /**
     * 构建增强的上下文消息
     */
    private List<ChatMessage> buildEnhancedContext(List<ChatMessageModel> originalMessages, RAGQueryResult queryResult) {
        List<ChatMessage> enhancedMessages = new ArrayList<>();

        // 如果有查询结果，添加系统提示
        if (queryResult.isSuccess() && StringUtils.isNotBlank(queryResult.getContext())) {
            String systemPrompt = String.format(
                "你是一个专业的博客助手，基于以下博客文章信息回答用户问题：\n\n%s\n\n" +
                "请基于上述信息回答用户问题，如果没有找到相关信息，请诚实地告诉用户。\n" +
                "回答要准确、有用，并尽可能引用具体的文章信息。",
                queryResult.getContext()
            );

            enhancedMessages.add(new ChatMessage("system", systemPrompt));
        }

        // 添加历史消息（保持原逻辑）
        List<ChatMessageModel> limitedMessages = originalMessages.size() > 4 ?
            originalMessages.subList(originalMessages.size() - 4, originalMessages.size()) :
            originalMessages;

        limitedMessages.forEach(chatMessageModel -> {
            enhancedMessages.add(new ChatMessage(
                chatMessageModel.getRole(),
                Html2TextUtil.getContent(chatMessageModel.getContent())
            ));
        });

        return enhancedMessages;
    }

    /**
     * 发送RAG增强的消息
     */
    private GlobalResult sendMessageWithRAG(User user, List<ChatMessage> messages, String model, RAGQueryResult queryResult) {
        if (StringUtils.isBlank(model)) {
            model = "deepseek-chat";
        }

        OpenAiService service = new OpenAiService(token, Duration.ofSeconds(180));
        ChatCompletionRequest completionRequest = ChatCompletionRequest.builder()
                .model(model)
                .stream(true)  // 使用流式响应
                .messages(messages)
                .build();

        // 发送查询信息给用户（可选）
        if (queryResult.isSuccess() && queryResult.getArticles() != null && !queryResult.getArticles().isEmpty()) {
            String queryInfo = String.format("📚 找到 %d 篇相关文章\n\n", queryResult.getArticles().size());
            try {
                // 尝试通过SSE发送查询信息
                sseService.send(user.getIdUser(), queryInfo);
            } catch (Exception e) {
                System.out.println(queryInfo);
            }
        }

        // 使用流式响应，通过SSE发送给前端
        service.streamChatCompletion(completionRequest)
                .doOnError(Throwable::printStackTrace)
                .blockingForEach(chunk -> {
                    if (chunk.getChoices().isEmpty() || chunk.getChoices().get(0).getMessage() == null) {
                        return;
                    }
                    String text = chunk.getChoices().get(0).getMessage().getContent();
                    if (text == null) {
                        return;
                    }
                    System.out.print(text);
                    try {
                        boolean sent = sseService.send(user.getIdUser(), text);
                        if (!sent) {
                            System.out.println("\n警告: 无法向用户 " + user.getIdUser() + " 发送消息，SSE连接可能已断开");
                        }
                    } catch (IllegalArgumentException e) {
                        System.out.println("\n警告: " + e.getMessage());
                        // 不抛出异常，继续处理后续消息
                    }
                });

        service.shutdownExecutor();
        return GlobalResultGenerator.genSuccessResult();
    }

    /**
     * 测试SuperSQL查询
     */
    @GetMapping("/test")
    public GlobalResult testQuery(@RequestParam String question) {
        RAGQueryResult result = superSQLService.queryByNaturalLanguage(question);
        return GlobalResultGenerator.genSuccessResult(result);
    }
}